# 若出现明明网页发出了文件，通过调试软件也看到了消息，而且用该调试软件发的消息开发板也能收到，但是开发板并没有收到网页发来的消息，可以尝试
# 排查以下问题：是否网页分包过大、频率过高。可以逐步缩小，直到能正确接收为止。我目前设置的是16B，1秒1发，仍然有丢包的风险，而且长时间传输
# 后就会直接罢工。解决方法是把服务质量换成QoS 2，反正理论上换模型这种事也不会经常发生，慢点就慢点吧
import image, time, os, sys, gc
from libs.PlatTasks import DetectionApp
from libs.PipeLine import PipeLine
from libs.Utils import *
from media.sensor import *
from media.display import *
from media.media import *
from ybUtils.mqtt import MQTTClient  # 导入MQTT客户端模块 / Import MQTT client module
import network  # 导入网络模块，用于WiFi连接 / Import network module for WiFi connection
import os      # 导入操作系统模块 / Import operating system module
import machine # 导入硬件控制模块 / Import hardware control module
import time    # 导入时间模块，用于延时操作 / Import time module for delays

# WiFi配置参数 / WiFi configuration parameters
WIFI_SSID = "box_of_Shittim"          # WiFi名称 / WiFi name
WIFI_PASSWORD = "1145141919810" # WiFi密码 / WiFi password

# MQTT服务器配置参数 / MQTT server configuration parameters
MQTT_BROKER = "47.109.75.129"  # MQTT服务器地址 / MQTT broker address
MQTT_PORT = 1883                # MQTT服务器端口 / MQTT broker port
MQTT_SEND_TOPIC = "image"    # MQTT主题 / MQTT topic
MQTT_RECEIVE_TOPIC = "finish"    # MQTT主题 / MQTT topic
MQTT_PEST_DETECTION_TOPIC = "pest/detection"
MQTT_MODEL_UPDATE_TOPIC = "model/update"
picture_width = 1280
picture_height = 720
DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

# Set display mode: options are 'hdmi', 'lcd', 'lt9611', 'st7701', 'hx8399'
# 'hdmi' defaults to 'lt9611' (1920x1080); 'lcd' defaults to 'st7701' (800x480)
display_mode = "lt9611"

# Define the input size for the RGB888P video frames
rgb888p_size = [1280, 720]

# Set root directory path for model and config
root_path = "/sdcard/mp_deployment_source/"

# Load deployment configuration
deploy_conf = read_json(root_path + "/deploy_config.json")
kmodel_path = root_path + deploy_conf["kmodel_path"]              # KModel path
labels = deploy_conf["categories"]                                # Label list
confidence_threshold = deploy_conf["confidence_threshold"]        # Confidence threshold
nms_threshold = deploy_conf["nms_threshold"]                      # NMS threshold
model_input_size = deploy_conf["img_size"]                        # Model input size
nms_option = deploy_conf["nms_option"]                            # NMS strategy
model_type = deploy_conf["model_type"]                            # Detection model type
anchors = []
if model_type == "AnchorBaseDet":
    anchors = deploy_conf["anchors"][0] + deploy_conf["anchors"][1] + deploy_conf["anchors"][2]

# Inference configuration
inference_mode = "video"                                          # Inference mode: 'video'
debug_mode = 0                                                    # Debug mode flag

model_update_in_progress = False
received_model_data = bytearray()
new_model_path = "/sdcard/mp_deployment_source/new_model.kmodel"
backup_model_path = "/sdcard/mp_deployment_source/backup_model.kmodel"


def connect_wifi(ssid, password):
    """
    连接WiFi并返回IP地址
    Connect to WiFi and return IP address

    参数 / Parameters:
    ssid: WiFi名称 / WiFi name
    password: WiFi密码 / WiFi password

    返回 / Returns:
    str: IP地址 / IP address
    """
    print(f"连接WIFI: {ssid}.. Connecting..")
    # 创建WiFi站点对象 / Create WiFi station object
    wifi_station = network.WLAN(0)
    # 连接到指定WiFi / Connect to specified WiFi
    wifi_station.connect(ssid, password)

    # 等待直到获取到IP地址 / Wait until IP address is obtained
    while wifi_station.ifconfig()[0] == '0.0.0.0':
        os.exitpoint()
    print("WIFI连接成功 WIFI Connected!")
    # 返回IP地址 / Return IP address
    return wifi_station.ifconfig()[0]

    def save_new_model():
        """保存新接收的模型文件"""
        global received_model_data, model_update_in_progress

        try:
            print("开始保存新模型文件...")

            # 备份当前模型
            if os.path.exists(kmodel_path):
                print("备份当前模型...")
                os.rename(kmodel_path, backup_model_path)

            # 保存新模型
            print(f"保存新模型到: {kmodel_path}")
            with open(kmodel_path, "wb") as f:
                f.write(received_model_data)

            # 验证文件大小
            saved_size = os.stat(kmodel_path)[6]  # 获取文件大小
            expected_size = len(received_model_data)

            if saved_size == expected_size:
                print(f"模型更新成功! 文件大小: {saved_size} 字节")

                # 删除备份文件
                if os.path.exists(backup_model_path):
                    os.remove(backup_model_path)

                # 重置状态
                received_model_data = bytearray()
                model_update_in_progress = False

                return True
            else:
                print(f"模型文件大小不匹配: 期望 {expected_size}, 实际 {saved_size}")
                # 恢复备份
                if os.path.exists(backup_model_path):
                    os.rename(backup_model_path, kmodel_path)
                return False

        except Exception as e:
            print(f"保存模型失败: {e}")
            # 恢复备份
            if os.path.exists(backup_model_path):
                os.rename(backup_model_path, kmodel_path)
            return False


def reload_model():
    """重新加载模型"""
    global det_app
    try:
        print("重新加载模型...")

        # 释放当前模型资源
        det_app.deinit()

        # 重新初始化检测应用
        det_app = DetectionApp(
            inference_mode, kmodel_path, labels, model_input_size,
            anchors, model_type, confidence_threshold, nms_threshold,
            rgb888p_size, display_size, debug_mode=debug_mode
        )

        # 重新配置预处理
        det_app.config_preprocess()

        print("模型重新加载成功!")
        return True

    except Exception as e:
        print(f"重新加载模型失败: {e}")
        return False


# 定义回调函数
def on_message(topic, msg):
    global flag, model_update_in_progress, received_model_data
    try:
        print("Topic: {}，Message: {}".format(topic, msg.decode()))
        if topic == b"finish" and msg.decode() == "continue":
            print("继续发送")
            flag = True  # 网页回复"continue"后继续开始发送
        elif topic == b"finish" and msg.decode() == "update over":
            print("收到模型更新完成信号，请重启设备")
            if model_update_in_progress and received_model_data:
                print(f"接收到的模型数据大小: {len(received_model_data)} 字节")
                flag = True

#                # 保存新模型
#                if save_new_model():
#                    # 重新加载模型
#                    if reload_model():
#                        print("模型更新并重新加载成功!")
#                    else:
#                        print("模型保存成功但重新加载失败")
#                else:
#                    print("模型保存失败")

            # 重置模型接收状态
            model_update_in_progress = False
            received_model_data = bytearray()

        elif topic == MQTT_MODEL_UPDATE_TOPIC.encode():
            if flag:
                flag = False
            print(f"收到模型数据包，大小: {len(msg)} 字节")

            # 开始或继续接收模型数据
            if not model_update_in_progress:
                model_update_in_progress = True
                received_model_data = bytearray()
                print(" 开始接收新模型...")

            # 追加数据
            received_model_data.extend(msg)
            print(f"已接收模型数据 {len(received_model_data) / len(msg)} 包、{len(received_model_data)} 字节")
    except:
        pass


# Create and initialize the video/display pipeline
pl = PipeLine(rgb888p_size=rgb888p_size, display_mode=display_mode)
sensor = Sensor(width=picture_width, height=picture_height)
sensor.reset()
sensor.set_hmirror(False)
sensor.set_vflip(False)
sensor.set_framesize(width=picture_width, height=picture_height, chn=CAM_CHN_ID_0)
sensor.set_pixformat(Sensor.RGB888, chn=CAM_CHN_ID_0)
pl.create(sensor)
display_size = pl.get_display_size()

# Initialize object detection application
det_app = DetectionApp(inference_mode,kmodel_path,labels,model_input_size,anchors,model_type,confidence_threshold,nms_threshold,rgb888p_size,display_size,debug_mode=debug_mode)

# Configure preprocessing for the model
det_app.config_preprocess()

# 初始化显示器
Display.init(Display.LT9611, width=DISPLAY_WIDTH, height=DISPLAY_HEIGHT, to_ide=True)

#MediaManager.init()
sensor.run()

block_length = 8192
wait_time = 5
flag = True
# 4G模块及摄像头初始化完成，进行图传任务
time.sleep(5)  # 等待摄像机稳定
# 主程序入口 / Main program entry
if __name__ == "__main__":
    # 连接WiFi网络 / Connect to WiFi network
    print(connect_wifi(WIFI_SSID, WIFI_PASSWORD))
    # 创建MQTT客户端实例 / Create MQTT client instance
    # 参数包括客户端ID、服务器地址和端口 / Parameters include client ID, broker address and port
    client = MQTTClient("YAHBOOM-K230", MQTT_BROKER, port=MQTT_PORT)
    # 连接到MQTT服务器 / Connect to MQTT broker
    client.connect()
    client.set_callback(on_message)
    client.subscribe(MQTT_SEND_TOPIC, qos=0)
    client.subscribe(MQTT_RECEIVE_TOPIC, qos=0)
    client.subscribe(MQTT_PEST_DETECTION_TOPIC, qos=0)
    client.subscribe(MQTT_MODEL_UPDATE_TOPIC, qos=2)
    print("subscribe success!")
    while True:
        if flag:

            with ScopedTiming("total", 1):
                img = pl.get_frame()                          # Capture current frame
                res = det_app.run(img)                        # Run inference
                if not type(res["idx"]) is list:
                    client.publish(MQTT_PEST_DETECTION_TOPIC, str(np.sum(res["idx"] != 0)))
                    det_app.draw_result(pl.osd_img, res)          # Draw detection results
                    pl.show_image()                               # Show result on display
                gc.collect()                                  # Run garbage collection

            # 捕获图像
            img = sensor.snapshot(chn=CAM_CHN_ID_0)
            # 显示原始图像
#            Display.show_image(img, x=0, y=0)
            img.draw_rectangle(0, 0, 60, 60)
            img = img.compress(quality=10)
            img.save("/data/tmp.jpg")
            with open('/data/tmp.jpg', 'rb') as img_bytes:
                img_bytes = img_bytes.read()
                block = int(len(img_bytes)/block_length)
                print(block)
                for i in range(block):
                    print(img_bytes[i*block_length:(i+1)*block_length])
                    client.publish(MQTT_SEND_TOPIC, img_bytes[i*block_length:(i+1)*block_length])
                    time.sleep_ms(wait_time)
                client.publish(MQTT_SEND_TOPIC, img_bytes[block*block_length:])  # 分包传输
                time.sleep_ms(wait_time)
            print("发送结束")
            flag = False
#            time.sleep(1)
        else:
            client.check_msg()
    # 断开与MQTT服务器的连接 / Disconnect from MQTT broker
    client.disconnect()
    # Cleanup: These lines will only run if the loop is interrupted (e.g., by an IDE break or external interruption)
    det_app.deinit()                                      # De-initialize detection app
    pl.destroy()                                          # Destroy pipeline instance
